import asyncio

from langchain_mcp_adapters.client import MultiServerMCPClient
from langgraph.prebuilt import create_react_agent

from agent.my_llm import llm

test_token = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJkZXYtdXNlciIsImlzcyI6Imh0dHBzOi8vd3d3Lmxhb3hpYW8uY29tIiwiaWF0IjoxNzU5OTk3NDk0LCJleHAiOjE3NjAwMDEwOTQsImF1ZCI6Im15LWRldi1zZXJ2ZXIiLCJzY29wZSI6Imxhb3hpYW8gaW52b2tlX3Rvb2xzIn0.J3WwbW1GCziaVhNesdLQG5BaiIpCFrCUCVdnMQ8L8-D79SXkls77A99dgIhugn5nssS1NgAgaWKPCbN5n99rhxpj_h_Z2mtKKKISGfk-DEz3cmkojJAy84UjrgRJkKxOkEW3kOI6ilmPHXmvpD7AJtGdPcmx-0cvsnme14csU9IgUKm5IStyd8Uq_L86sX2Mv5K2mnDWHicx12qpVrIBscSN3Ehg4iqlweFSyOkejTRxdPUtXKZCckxLcWhGDc2yFcCh9ZaoXL1bPHDQTT7BFbD2gVqAmhClpILqa0HVdpjO4R9hohPvj7nFu_5iU1ZAwN17UF84GhPnqabwcXCBig"

# Python MCP 服务端的连接配置
python_mcp_server_config = {
    'url': 'http://127.0.0.1:8080/streamable',
    'transport': 'streamable_http',
    'headers': {
        'Authorization': f'Bearer {test_token}',
    },
    # 'url': 'http://127.0.0.1:8080/sse',
    # 'transport': 'sse',
}


mcp_client = MultiServerMCPClient(
    {
        'python_mcp': python_mcp_server_config
    }
)


async def create_agent():
    """必须是异步函数"""
    mcp_tools = await mcp_client.get_tools()
    print(mcp_tools)

    # 查询mcp服务端定义的prompt
    p = await mcp_client.get_prompt(server_name='python_mcp', prompt_name='ask_about_topic', arguments={'topic': '大模型开发'})
    print(p)

    # 查询mcp服务端定义的resource
    data = await mcp_client.get_resources(server_name='python_mcp', uris='resource://config')
    print(data)     # [Blob 2041768728992]
    print(type(data))   # <class 'list'>
    if isinstance(data, list) and data:
        print(data[0])
        # metadata={'uri': 'resource://config'} data='{"theme":"dark","version":"1.2.0","features":["tools",
        # "resources"]}' mimetype='text/plain'
        print(data[0].data)
        # {"theme":"dark","version":"1.2.0","features":["tools","resources"]}

    agent = create_react_agent(
        llm,
        tools=mcp_tools,
        prompt='你是一个智能助手，尽可能的调用工具回答用户的问题'
    )

    ret = await agent.ainvoke(
        {"messages": [{"role": "user", "content": "今天，北京的天气怎么样？"}]}
    )
    print(ret['messages'])
    print(ret['messages'][-1].content)


if __name__ == '__main__':
    asyncio.run(create_agent())


